# coding=utf8
from __future__ import division
import os
import torch
import torch.utils.data as data
import PIL.Image as Image
from dataset.data_aug import *
import cv2
import numpy as np
import matplotlib.pyplot as plt
class dataset_test(data.Dataset):                     # user-defined  dataset framework
    def __init__(self, anno_pd, transforms=None):
        self.paths = anno_pd['image_id'].tolist()    # ndarray.tolist()  :  return the array as a list
        self.labels = anno_pd['classes'].tolist()        #anno_pd['label'] :  dict
        self.transforms = transforms

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, item):
        img_path = self.paths[item]
        # img = self.pil_loader(img_path)        # 1. read one data from file
        img =cv2.imread(img_path)                # 2. read one data :  yige data
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)   # [h,w,3]  RGB   preprocess
        # print(img.shape)
        if self.transforms is not None:
            img = self.transforms(img)           # return a data pair(image and label)
        label = self.labels[item]
        # return img, label
        # plt.imshow(img)
        # plt.show()
        return torch.from_numpy(img).float(), label
class dataset_train(data.Dataset):                     # user-defined  dataset framework
    def __init__(self, anno_pd, transforms, resize):
        self.paths = anno_pd['image_path'].tolist()    # ndarray.tolist()  :  return the array as a list
        self.labels = anno_pd['label'].tolist()        #anno_pd['label'] :  dict
        self.transforms = transforms
        self.resize = resize

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, item):
        img_path = self.paths[item]
        # img = self.pil_loader(img_path)
        # 1. read one data from file
        # print(img_path)
        image =cv2.imread(img_path)                # 2. read one data :  yige data
        image_ = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)   # [h,w,3]  RGB   preprocess
        min_scale=min(self.resize/image_.shape[0],self.resize/image_.shape[1])
        white_img=np.zeros((self.resize,self.resize,3),dtype=np.uint8)+255
        img_ = cv2.resize(image_, (0,0),fx=min_scale,fy=min_scale)        # todo
        real_img=white_img.copy()
        real_img[:img_.shape[0],:img_.shape[1],:]=img_
        img=real_img

        if self.transforms is not None:
            img = self.transforms(img)           # return a data pair(image and label)
        label = self.labels[item]
        # return img, label
        # plt.imshow(img)
        # plt.show()
        return torch.from_numpy(img).float(), label
class dataset(data.Dataset):
    def __init__(self, imgroot, anno_pd, transforms=None):
        self.root_path = imgroot
        self.paths = anno_pd['image_id'].tolist()
        self.labels = anno_pd['disease_class'].tolist()
        self.transforms = transforms

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, item):
        img_path = self.paths[item]
        # img = self.pil_loader(img_path)
        img =cv2.imread(img_path)
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)   # [h,w,3]  RGB

        if self.transforms is not None:
            img = self.transforms(img)
        label = self.labels[item]
        # return img, label
        return torch.from_numpy(img).float(), label

    def pil_loader(self,imgpath):
        with open(imgpath, 'rb') as f:
            with Image.open(f) as img:
                return img.convert('RGB')
class dataset_testaug(data.Dataset):
    def __init__(self, anno_pd, transforms=None):
        self.paths = anno_pd['image_id'].tolist()  # ndarray.tolist()  :  return the array as a list
        self.labels = anno_pd['classes'].tolist()  # anno_pd['label'] :  dict
        self.transforms = transforms

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, item):
        img_path = self.paths[item]
        # img = self.pil_loader(img_path)
        img =cv2.imread(img_path)
        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)   # [h,w,3]  RGB
        # img = cv2.flip(img,1)
        if self.transforms is not None:
            img= self.transforms(img)

        image_list = RandomTwelveCrop(size=(272,272))(img)
        image_list = Normalize_test([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image_list)
        label = self.labels[item]
        # return img, label
        image_list=np.array(image_list)
        return torch.from_numpy(image_list).float(), label

    def pil_loader(self,imgpath):
        with open(imgpath, 'rb') as f:
            with Image.open(f) as img:
                return img.convert('RGB')


def collate_fn(batch):
    imgs = []
    label = []

    for sample in batch:
        imgs.append(sample[0])
        label.append(sample[1])
        # print(sample[0].shape, sample[1])
    return torch.stack(imgs, 0), label


def change(item):
    if item == 0:
        item = "DESERT"
    elif item == 1:
        item = "OCEAN"
    elif item == 2:
        item = "MOUNTAIN"
    elif item == 3:
        item = "FARMLAND"
    elif item == 4:
        item = "LAKE"
    else:
        item = "CITY"
    return item
def rechange(item):
    if item == "DESERT":
        item = "0"
    elif item == "OCEAN":
        item = "1"
    elif item == "MOUNTAIN":
        item = "2"
    elif item == "FARMLAND":
        item = "3"
    elif item == "LAKE":
        item = "4"
    else:
        item = "5"
    return item
if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import pandas as pd
    rawdata_root = '/media/hszc/data/detao/data/baidu/datasets'
    train_pd = pd.read_csv("/media/hszc/data/detao/data/baidu/datasets/train.txt", sep=" ",
                           header=None, names=['ImageName', 'label'])
    val_pd = pd.read_csv("/media/hszc/data/detao/data/baidu/datasets/test_answer.txt", sep=" ",
                         header=None, names=['ImageName', "label"])
    val_pd['ImageName'] = val_pd['ImageName'].apply(lambda x: os.path.join('test', x))
    train_pd['ImageName'] = train_pd['ImageName'].apply(lambda x: os.path.join('train', x))
    data_transforms = {
        'train': Compose([
            RandomRotate(angles=(-15,15)),
            ExpandBorder(size=(368, 368), resize=True),
            RandomResizedCrop(size=(336, 336)),
            # RandomHflip(),
            # transforms.RandomResizedCrop(336,scale=(0.49,1.0)),
            # transforms.RandomHorizontalFlip(),
            # transforms.ToTensor(),   # 0-255 to 0-1
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]),
        'val': Compose([
            ExpandBorder(size=(368, 368), resize=True),
            # CenterCrop(size=(336, 336)),
            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    data_set = {}
    data_set['train'] = dataset(imgroot=rawdata_root, anno_pd=train_pd,
                                transforms=data_transforms["train"],
                                )
    data_set['val'] = dataset(imgroot=rawdata_root, anno_pd=val_pd,
                              transforms=data_transforms["val"],
                              )
    for data in data_set["val"]:
        print("22222")